import os
import threading
import time
import random
from datetime import datetime

# 配置参数（请根据实际情况修改）
TEST_DIR = "H:\\test_concurrent"  # U盘测试目录
FILE_SIZE_MB = 100  # 单个文件大小(MB)
WRITE_THREAD_COUNT = 2  # 写线程数量
READ_THREAD_COUNT = 2  # 读线程数量
BLOCK_SIZE = 4 * 1024 * 1024  # 4MB块大小（匹配U盘常见块大小）
SAMPLE_INTERVAL = 0.2  # 速度采样间隔（秒）

# 全局变量
write_results = []  # 写线程结果：(线程ID, 耗时, 平均速度, 峰值速度, 速度列表)
read_results = []  # 读线程结果：(线程ID, 耗时, 平均速度, 峰值速度, 速度列表)
stop_event = threading.Event()  # 线程终止信号


def generate_random_data(size):
    """生成随机数据（避免系统缓存优化）"""
    return os.urandom(size)


def write_worker(thread_id):
    """写线程工作函数（仅操作自己的专属文件）"""
    filename = f"{TEST_DIR}\\write_file_{thread_id}.dat"
    speed_samples = []  # 记录每个采样间隔的速度
    total_written = 0  # 总写入字节数

    try:
        # 打开文件（每次写入都是新文件，避免覆盖冲突）
        with open(filename, "wb") as f:
            start_time = time.time()
            last_sample_time = start_time
            last_written = 0
            data = generate_random_data(BLOCK_SIZE)  # 预生成数据块

            # 循环写入直到达到目标大小
            while total_written < FILE_SIZE_MB * 1024 * 1024:
                if stop_event.is_set():
                    raise Exception("测试被终止")

                # 计算本次写入大小（最后一块可能不足BLOCK_SIZE）
                write_size = min(BLOCK_SIZE, FILE_SIZE_MB * 1024 * 1024 - total_written)
                f.write(data[:write_size])  # 写入数据
                total_written += write_size

                # 按间隔采样速度
                current_time = time.time()
                if current_time - last_sample_time >= SAMPLE_INTERVAL:
                    elapsed = current_time - last_sample_time
                    written_mb = (total_written - last_written) / (1024 * 1024)  # 转换为MB
                    speed = written_mb / elapsed
                    speed_samples.append(speed)

                    # 更新采样点
                    last_sample_time = current_time
                    last_written = total_written

        # 计算线程结果
        end_time = time.time()
        total_duration = end_time - start_time
        avg_speed = FILE_SIZE_MB / total_duration if total_duration > 0 else 0
        peak_speed = max(speed_samples) if speed_samples else 0

        write_results.append((thread_id, total_duration, avg_speed, peak_speed, speed_samples))
        print(f"[写线程{thread_id}] 完成！平均: {avg_speed:.2f}MB/s，峰值: {peak_speed:.2f}MB/s")

    except Exception as e:
        print(f"[写线程{thread_id}] 出错: {str(e)}")
        stop_event.set()  # 触发其他线程终止


def read_worker(thread_id):
    """读线程工作函数（仅操作自己的专属文件）"""
    filename = f"{TEST_DIR}\\read_file_{thread_id}.dat"
    speed_samples = []  # 记录每个采样间隔的速度
    total_read = 0  # 总读取字节数

    try:
        # 检查文件是否存在
        if not os.path.exists(filename):
            raise Exception(f"文件不存在: {filename}")

        # 打开文件读取
        with open(filename, "rb") as f:
            start_time = time.time()
            last_sample_time = start_time
            last_read = 0

            # 循环读取直到达到目标大小
            while total_read < FILE_SIZE_MB * 1024 * 1024:
                if stop_event.is_set():
                    raise Exception("测试被终止")

                # 读取数据块
                data = f.read(BLOCK_SIZE)
                if not data:  # 文件提前结束
                    raise Exception(f"文件不完整，仅读取{total_read / (1024 * 1024):.2f}MB")
                total_read += len(data)

                # 按间隔采样速度
                current_time = time.time()
                if current_time - last_sample_time >= SAMPLE_INTERVAL:
                    elapsed = current_time - last_sample_time
                    read_mb = (total_read - last_read) / (1024 * 1024)  # 转换为MB
                    speed = read_mb / elapsed
                    speed_samples.append(speed)

                    # 更新采样点
                    last_sample_time = current_time
                    last_read = total_read

        # 计算线程结果
        end_time = time.time()
        total_duration = end_time - start_time
        avg_speed = FILE_SIZE_MB / total_duration if total_duration > 0 else 0
        peak_speed = max(speed_samples) if speed_samples else 0

        read_results.append((thread_id, total_duration, avg_speed, peak_speed, speed_samples))
        print(f"[读线程{thread_id}] 完成！平均: {avg_speed:.2f}MB/s，峰值: {peak_speed:.2f}MB/s")

    except Exception as e:
        print(f"[读线程{thread_id}] 出错: {str(e)}")
        stop_event.set()  # 触发其他线程终止


def main():
    # 初始化测试目录
    try:
        os.makedirs(TEST_DIR, exist_ok=True)
        print(f"测试目录: {TEST_DIR}")
        print(f"配置: {WRITE_THREAD_COUNT}个写线程 + {READ_THREAD_COUNT}个读线程，单个文件{FILE_SIZE_MB}MB")
    except Exception as e:
        print(f"初始化目录失败: {str(e)}")
        return

    # 创建并启动所有线程
    threads = []

    # 启动写线程
    for i in range(WRITE_THREAD_COUNT):
        t = threading.Thread(target=write_worker, args=(i,), name=f"Write-{i}")
        threads.append(t)
        t.start()

    # 启动读线程
    for i in range(READ_THREAD_COUNT):
        t = threading.Thread(target=read_worker, args=(i,), name=f"Read-{i}")
        threads.append(t)
        t.start()

    # 等待所有线程完成
    print("\n所有线程启动完成，开始并发测试...")
    for t in threads:
        t.join()

    # 输出汇总结果
    print("\n" + "=" * 60)
    print(f"测试完成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

    print("\n【注意事项】")
    print("1. 结果反映U盘在并发读写时的真实性能（无锁竞争）")
    print("2. 峰值速度受U盘缓存影响，平均速度反映持续性能")
    print("3. 测试文件可手动删除: " + TEST_DIR)
    print("=" * 60)


if __name__ == "__main__":
    main()
